## 从视频文件中获取图片帧
## 将指定文件夹中的所有.jpg图片按顺序重命名，新的名称从你指定的数字开始
import cv2
import os

def extract_frames(
    video_path, output_folder, frames_per_second=2, interval_seconds=None, saved_frame_count=2206
):
    # 创建输出文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
        print(f"创建输出文件夹: {output_folder}")
    else:
        print(f"输出文件夹已存在: {output_folder}")

    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件: {video_path}")
        return

    # 获取视频的帧率和总帧数
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps if fps else 0
    print(f"视频帧率 (FPS): {fps}")
    print(f"总帧数: {total_frames}")
    print(f"视频时长 (秒): {duration:.2f}")

    if interval_seconds:
        # 如果指定了时间间隔，计算对应的帧间隔
        frame_interval = int(fps * interval_seconds)
        print(f"每隔 {interval_seconds} 秒提取一帧，约等于每隔 {frame_interval} 帧提取一帧。")
    else:
        # 否则按照每秒提取的帧数计算帧间隔
        frame_interval = int(fps // frames_per_second) if fps else 0
        if frame_interval <= 0:
            print(f"警告: frames_per_second ({frames_per_second}) 大于视频帧率 ({fps}). 将每隔1帧提取一帧。")
            frame_interval = 1
        print(f"每隔 {frame_interval} 帧提取一帧。")

    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            print("视频读取完毕或发生错误。")
            break

        # 打印当前处理的帧数
        if frame_count % 100 == 0:
            print(
                f"当前处理帧数: {frame_count}/{total_frames} ({(frame_count/total_frames)*100:.2f}%)"
            )

        if frame_count % frame_interval == 0:
            # 保存帧
            output_path = os.path.join(
                output_folder, f"frame_{saved_frame_count:05d}.jpg"
            )
            success = cv2.imwrite(output_path, frame)
            if success:
                print(f"保存帧 {saved_frame_count} 到 {output_path}")
                saved_frame_count += 1
            else:
                print(f"保存帧 {saved_frame_count} 失败。")

        frame_count += 1

    cap.release()
    print(f"提取完成，已保存 {saved_frame_count} 帧到 {output_folder}.")

def rename_images(folder_path, start_number):
    # 获取文件夹中的所有文件
    files = sorted([f for f in os.listdir(folder_path) if f.lower().endswith('.jpg')])

    # 遍历文件并重命名
    for i, filename in enumerate(files, start=start_number):
        # 构造新的文件名
        new_name = f"frame_{i:05d}.jpg"  # 保持5位数的格式
        old_path = os.path.join(folder_path, filename)
        new_path = os.path.join(folder_path, new_name)
        
        # 重命名文件
        os.rename(old_path, new_path)
        print(f"Renamed {filename} to {new_name}")

# 示例用法
video_file = "/home/shuai/Videos/dataset_20250227.mp4"  # 替换为你的输入视频文件路径
output_directory = "/home/shuai/Videos/images"  # 替换为你想要保存帧的文件夹

## frames_per_second表示每秒提取多少帧，interval_seconds表示每隔多少秒提取一帧，saved_frame_count表示保存的帧图片名称从什么数字开始
extract_frames(video_file, output_directory, frames_per_second=None, interval_seconds=1, saved_frame_count=10811)

## 将指定文件夹中的所有.jpg图片按顺序重命名，新的名称从你指定的数字开始
# start_number = 1080
# rename_images(output_directory, start_number)
